import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text
from sklearn.metrics import accuracy_score, classification_report
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.compose import ColumnTransformer

# 1. 加载数据（假设数据已整理为CSV，去除无关列如thedate、MovieID）
data = pd.read_csv('movie_data.csv')
data = data.dropna()  # 去除缺失值

# 2. 目标变量：将票房转换为二分类（高于中位数为高票房，否则为低票房）
data['BoxOffice_class'] = data['BoxOffice'].apply(
    lambda x: 1 if x > data['BoxOffice'].median() else 0
)

# 3. 特征工程：处理分类变量CityLevel（独热编码）
features = data[['CityLevel', 'AudienceCount', 'ShowCount']]
target = data['BoxOffice_class']

# 4. 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    features, target, test_size=0.3, random_state=42
)

# 5. 数据预处理：对CityLevel进行独热编码，数值型特征标准化
preprocessor = ColumnTransformer(
    transformers=[
        ('cat', OneHotEncoder(), ['CityLevel']),
        ('num', StandardScaler(), ['AudienceCount', 'ShowCount'])
    ])
X_train_pre = preprocessor.fit_transform(X_train)
X_test_pre = preprocessor.transform(X_test)


# 模型1：C5.0等价实现（使用信息增益，类似C4.5）
model_c5 = DecisionTreeClassifier(
    criterion='entropy',  # 信息增益
    max_depth=5,
    min_samples_split=10,
    random_state=42
)

# 模型2：C&R Tree（使用Gini指数，sklearn默认）
model_cr = DecisionTreeClassifier(
    criterion='gini',  # Gini指数
    max_depth=5,
    min_samples_split=10,
    random_state=42
)

# 训练模型
model_c5.fit(X_train_pre, y_train)
model_cr.fit(X_train_pre, y_train)

# 预测
y_pred_c5 = model_c5.predict(X_test_pre)
y_pred_cr = model_cr.predict(X_test_pre)


# 评估函数
def evaluate_model(y_test, y_pred, model_name):
    accuracy = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, target_names=['低票房', '高票房'])
    print(f"=== {model_name} 评估结果 ===")
    print(f"准确率：{accuracy:.4f}")
    print("分类报告：\n", report)

# 评估结果
evaluate_model(y_test, y_pred_c5, "C5.0（信息增益）")
evaluate_model(y_test, y_pred_cr, "C&R Tree（Gini指数）")

